import matplotlib.pyplot as plt
from pathlib import Path
import numpy as np

# 类别名称列表
# classes = ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle', 'bus', 'car', 'cat', 'chair', 'cow', 'diningtable', 'dog',
#            'horse', 'motorbike', 'person', 'pottedplant', 'sheep', 'sofa', 'train', 'tvmonitor']
classes = ['coal_gangue', 'coal']
colors = plt.cm.hsv(np.linspace(0, 1, len(classes))).tolist()

# 设定标签文件夹的路径
label_dir = Path(r'F:\Dataset\origin_data\labels\trainval')

# 读取所有标签文件
labels = []
for label_file in label_dir.glob('*.txt'):
    with open(label_file) as f:
        for line in f:
            class_id = int(line.split()[0])
            labels.append(classes[class_id])  # 使用类别名称

# 统计每个类别的频率
labels, counts = np.unique(labels, return_counts=True)

# 绘制直方图
plt.figure(figsize=(25, 15))

# plt.bar(labels, counts)

bars = plt.bar(classes, counts)

# 为每个条形分配颜色
for bar, color in zip(bars, colors):
    bar.set_color(color)

plt.xlabel('Class',fontdict={'fontsize':20})
plt.ylabel('Instances',fontdict={'fontsize':20})
plt.title('Label Distribution in Pascal VOC2012',fontdict={'fontsize':30})
# plt.xticks(rotation=90)  # 将类别名称旋转以便阅读
plt.xticks(ticks=np.arange(len(classes)), labels=classes, rotation=90, fontsize=12)
# 保存直方图
plt.savefig('OriginData.jpg',dpi=300)
# plt.show()
